import os
import cv2
import copy
import random
import xml.dom.minidom
from tqdm import tqdm
import numpy as np


def ReadXml(FilePath):
    """输入一个xml文件路径， 将xml文件内所有object目标下的信息返回[[xmin, ymin, xmax , ymax, name], ...]"""
    if os.path.exists(FilePath) is False:
        return None
    dom = xml.dom.minidom.parse(FilePath)
    # 得到文档元素对象
    root_ = dom.documentElement

    object_ = root_.getElementsByTagName('object')
    info = []
    # filename = root.getElementsByTagName('filename')[0].firstChild.data
    for object_1 in object_:
        name = object_1.getElementsByTagName("name")[0].firstChild.data
        bndbox = object_1.getElementsByTagName("bndbox")[0]
        xmin = int(bndbox.getElementsByTagName("xmin")[0].firstChild.data)
        ymin = int(bndbox.getElementsByTagName("ymin")[0].firstChild.data)
        xmax = int(bndbox.getElementsByTagName("xmax")[0].firstChild.data)
        ymax = int(bndbox.getElementsByTagName("ymax")[0].firstChild.data)
        # if name=='40\'':
        #   print(filename,xmin,ymin,xmax,ymax)
        # if 'limitRate'==name or 'limit'==name:
        info.append([xmin, ymin, xmax, ymax, name])
    return info


def WriteXml(infos, W, H):
    """
     写图片的xml文件

    :param infos: n维list,第一维是要写的xml文件的路径,  第2到n维是[xmin, ymin, xmax , ymax, name]格式数据
    :param W:  W是xml文件所描述图像的宽
    :param H:  H是xml文件所描述图像的高
    """
    FileName = infos[0].split("/")[-1]
    xml_file = open((infos[0]), 'w')
    xml_file.write('<annotation>\n')
    xml_file.write('    <folder>VOC2007</folder>\n')
    xml_file.write('    <filename>' + FileName + '</filename>\n')
    xml_file.write('    <size>\n')
    xml_file.write('        <width>' + str(W) + '</width>\n')
    xml_file.write('        <height>' + str(H) + '</height>\n')
    xml_file.write('        <depth>3</depth>\n')
    xml_file.write('    </size>\n')
    for img_each_label in infos[1:]:
        spt = img_each_label[4]
        xml_file.write('    <object>\n')
        xml_file.write('        <name>' + str(spt) + '</name>\n')
        xml_file.write('        <pose>Unspecified</pose>\n')
        xml_file.write('        <truncated>0</truncated>\n')
        xml_file.write('        <difficult>0</difficult>\n')
        xml_file.write('        <bndbox>\n')
        xml_file.write('            <xmin>' + str(img_each_label[0]) + '</xmin>\n')
        xml_file.write('            <ymin>' + str(img_each_label[1]) + '</ymin>\n')
        xml_file.write('            <xmax>' + str(img_each_label[2]) + '</xmax>\n')
        xml_file.write('            <ymax>' + str(img_each_label[3]) + '</ymax>\n')
        xml_file.write('        </bndbox>\n')
        xml_file.write('    </object>\n')

    xml_file.write('</annotation>')


def CountLabelKind(Path):
    """
    统计文件夹及其子文件夹下所有标签种类及其数量

    :param Path: 待统计文件夹
    :return: { ”LabelName":num,...}
    """
    LabelDict = {}
    print("Star to count label kinds....")
    for root, dirs, files in os.walk(Path):
        for file in tqdm(files):
            if file[-1] == 'l':
                Infos = ReadXml(root + "/" + file)
                for Info in Infos:
                    if Info[-1] not in LabelDict.keys():
                        LabelDict[Info[-1]] = 1
                    else:
                        LabelDict[Info[-1]] += 1

    return dict(sorted(LabelDict.items(), key=lambda x: x[0]))


def RenameJpgAndXml(RenamePath):
    """
    将文件夹下jpg文件按顺序重命名，如果相应的xml文件存在则将其也同步重命名

    :param RenamePath: 待重命名文件夹路径
    """
    print("Rename files.........")
    name = 0
    os.chdir(RenamePath)
    DirAndFile = os.listdir(RenamePath)
    for file in tqdm(DirAndFile):
        if file[-1] == 'g':
            name += 1
            os.rename(file, "AS%05d.jpg" % name)  # rename jpg
            if os.path.exists(file[:-3] + "xml") is True:
                os.rename(file[:-3] + "xml", "AS%05d.xml" % name)  # rename xml


def ResizeJpgAndXml(OldJpg, OldInfos, NewW, NewH):
    """
    输入一张图片和它所有目标的Info，将图片resize成(NewW, NewH)

    :param OldJpg:输入的图片
    :param OldInfos:输入的Infos
    :param NewW:输出图片的W
    :param NewH:输出图片的H
    :return:新图片， 新Infos
    """
    OldH, OldW, _ = OldJpg.shape
    OldIndex = np.array(OldInfos)[:, :-1].astype(float).astype(int)
    MultiplicationCoefficient = np.array([NewW / OldW, NewH / OldH, NewW / OldW, NewH / OldH])
    NewIndex = ((OldIndex * MultiplicationCoefficient).astype(int) + np.array([0, 0, 1, 1]))
    WH = [NewW, NewH]
    for i, column in enumerate((2, 3)):
        NewIndex[:, column][NewIndex[:, column] > WH[i]] = WH[i]
    NewInfos = NewIndex.tolist()
    Result = np.array(OldInfos)[:, -1].tolist()
    for i in range(len(NewInfos)):
        NewInfos[i].append(Result[i])
    return cv2.resize(OldJpg, (NewW, NewH)), NewInfos


def CutLabelWithBackground(Img, IndexInSmall, CutRateRange, XMin, YMin, XMax, YMax):
    """
    利用BigImg坐标在SmallImg上随机裁剪出包含BigImg里目标的图片

    :param Img:SmallImg
    :param IndexInSmall:BigImg里各目标在SmallImg里的坐标， 不含name
    :param CutRateRange:背景占图片比例
    :param XMin:BigImg在SmallImg的位置
    :param YMin:
    :param XMax:
    :param YMax:
    :return:裁剪后的图片， 裁剪后图片目标的坐标， 坐标越界数量
    """
    XMin, YMin, XMax, YMax = int(XMin), int(YMin), int(XMax), int(YMax)
    CutH = random.randint(int((YMax - YMin) * CutRateRange[0]), int((YMax - YMin) * CutRateRange[1])) + YMax - YMin
    CutW = random.randint(int((XMax - XMin) * CutRateRange[0]), int((XMax - XMin) * CutRateRange[1])) + XMax - XMin
    if CutH > (YMax - YMin):  # H比原H大
        RandH = random.randint(0, CutH - (YMax - YMin))
    else:
        RandH = random.randint(CutH - (YMax - YMin), 0)
    if CutW > XMax - XMin:  # W比原W大
        RandW = random.randint(0, CutW - (XMax - XMin))
    else:
        RandW = random.randint(CutW - (XMax - XMin), 0)

    while RandH > YMin or Img.shape[0] < CutH - RandH + YMin or RandW > XMin or Img.shape[1] < CutW - RandW + XMin:
        if RandH > YMin:
            RandH = random.randint(0, YMin)
        if Img.shape[0] < CutH - RandH + YMin:
            RandH = random.randint(RandH, CutH - (YMax - YMin))
        if RandW > XMin:
            RandW = random.randint(0, XMin)
        if Img.shape[1] < CutW - RandW + XMin:
            RandW = random.randint(RandW, CutW - (XMax - XMin))
    OutPutImg = Img[YMin - RandH:YMin - RandH + CutH, XMin - RandW:XMin - RandW + CutW]
    OutPutOffset = np.array([XMin - RandW, YMin - RandH, XMin - RandW, YMin - RandH]).astype(int)
    OutPutIndex = (IndexInSmall - OutPutOffset)
    AxisOutBoundaryNum = len(OutPutIndex[OutPutIndex < 0])
    return OutPutImg, OutPutIndex, AxisOutBoundaryNum


def RemoveFilesByKeyWord(Path, KeyWord):
    """
    删除路径下文件名带KeyWord的文件

    :param Path: 文件夹路径
    :param KeyWord: 需要删除的KeyWord
    """
    cnt = 0
    for root, dirs, files in os.walk(Path):
        for file in tqdm(files):
            if KeyWord in file:
                os.remove(root + "\\" + file)
                cnt += 1
    print("Remove %d files" % cnt)


def RemoveFileByName(Path, XmlList):
    """
    将路径下所有在XmlList里的xml及其对应jpg文件删除

    :param Path: 文件夹路径
    :param XmlList: 需要删除的xml文件列表
    :return:
    """
    print("Starting remove files...")
    for root, dirs, files in os.walk(Path):
        for file in tqdm(files):
            if file in XmlList:
                os.remove(root + "\\" + file)
                os.remove(root + "\\" + file[:-3] + "jpg")


def FindSmallLabel(BigLabelNamePath, SmallPath):
    """
    查找BigLabel文件对应的SmallLabel文件

    :param BigLabelNamePath: BigLabel 的绝对路径
    :param SmallPath: SmallLabel 的绝对路径
    :return:返回对应的SmallLabel文件名或None
    """
    BigFileName = BigLabelNamePath.split("\\")[-1]
    BigFileNameElement = BigFileName.split("_")
    SuspectedFiles = []
    TrueFile = None
    for root, dirs, files in os.walk(SmallPath):
        for file in files:
            if file[-1] == "l":
                if len(file) >= 20:
                    FilePrefix = file[-9:-4]
                else:
                    FilePrefix = file[:-4]
                if FilePrefix == BigFileNameElement[0]:
                    SuspectedFiles.append(file)
    # print(SuspectedFiles)
    for SuspectedFile in SuspectedFiles:
        Infos = ReadXml(SmallPath + "\\" + SuspectedFile)
        for Info in Infos:
            if str(Info[0]) == BigFileNameElement[1] and str(Info[1]) == BigFileNameElement[2] and str(Info[2]) == \
                    BigFileNameElement[3] and str(Info[3]) == BigFileNameElement[4]:
                TrueFile = SuspectedFile[:-3] + "jpg"
                return TrueFile
    return TrueFile


def CutInfos(XMin, YMin, XMax, YMax, NewInfo0, OldInfos, CoverRate=0.1):
    """
    将(XMin, YMin)作为新起点， 返回新起点后的Infos

    :param XMin: 新Img所在原Img上的XMin
    :param YMin: 新Img所在原Img上的YMin
    :param XMax: 新Img所在原Img上的XMax
    :param YMax: 新Img所在原Img上的YMax
    :param NewInfo0: 新Infos的Infos[0],即xml路径
    :param OldInfos: 原图的Infos
    :param CoverRate: 标签被允许的最大覆盖率
    :return: 新图对应的Infos
    """
    NewInfos = [NewInfo0]
    OldInfosCopy = copy.deepcopy(OldInfos)
    for OldInfo in OldInfosCopy:
        CoverX = (OldInfo[2] - OldInfo[0]) * CoverRate
        CoverY = (OldInfo[3] - OldInfo[1]) * CoverRate * 2
        if OldInfo[0] >= XMin - CoverX and OldInfo[1] >= YMin - CoverY and OldInfo[2] <= XMax + CoverX and OldInfo[
            3] <= YMax + CoverY:
            if OldInfo[0] < XMin:
                OldInfo[0] = XMin
            if OldInfo[1] < YMin:
                OldInfo[1] = YMin
            if OldInfo[2] > XMax:
                OldInfo[2] = XMax
            if OldInfo[3] > YMax:
                OldInfo[3] = YMax
            NewInfo = [OldInfo[0] - XMin, OldInfo[1] - YMin, OldInfo[2] - XMin, OldInfo[3] - YMin, OldInfo[4]]
            NewInfos.append(NewInfo)
    return NewInfos


def PutInfosRand300300(SecImg, Infos):
    H, W = 300, 300
    for Info in Infos:
        img = np.random.randint(0, 255, (H, W, 3)).astype(np.uint8)
        RandW = random.randint(0, W - Info[2] + Info[0])
        RandH = random.randint(0, H - Info[3] + Info[1])
        img[RandH:RandH + Info[3] - Info[1], RandW:RandW + Info[2] - Info[0]] = SecImg[Info[1]:Info[3], Info[0]:Info[2]]
        cv2.imshow("test", img)
        cv2.waitKey(0)


def BrightnessAverage(BigPic, SmallPic):
    BigYuv = cv2.cvtColor(BigPic, cv2.COLOR_BGR2YUV)
    SmallYuv = cv2.cvtColor(SmallPic, cv2.COLOR_BGR2YUV)
    BigAverBri, _, _, _ = cv2.mean(BigYuv)
    SmallAverBri, _, _, _ = cv2.mean(SmallYuv)
    Revise = SmallAverBri - BigAverBri
    SmallYuv[:, :, 0] = cv2.subtract(SmallYuv[:, :, 0], Revise)
    return cv2.cvtColor(SmallYuv, cv2.COLOR_YUV2BGR)


def FindNumberEdge(img, BackGroundColor):
    H, W, C = img.shape
    XMin, YMin, XMax, YMax = 0, 0, 0, 0
    for i in range(W):
        for channel in range(C):
            if False in (img[:, i, channel] == BackGroundColor[channel]):
                # XMin = i - 1
                XMin = i
                break
        else:
            continue
        break
    for i in range(W):
        for channel in range(C):
            if False in (img[:, W - 1 - i, channel] == BackGroundColor[channel]):
                # XMax = W - i + 1
                XMax = W - i
                break
        else:
            continue
        break
    for i in range(H):
        for channel in range(C):
            if False in (img[i, :, channel] == BackGroundColor[channel]):
                # YMin = i - 1
                YMin = i
                break
        else:
            continue
        break
    for i in range(H):
        for channel in range(C):
            if False in (img[H - 1 - i, :, channel] == BackGroundColor[channel]):
                # YMax = H - i + 1
                YMax = H - i
                break
        else:
            continue
        break
    if XMin < 0:
        XMin = 0
    if YMin < 0:
        YMin = 0
    return XMin, YMin, XMax, YMax


def FindNumberEdge2(img, color_text):
    bgr = (color_text[2], color_text[1], color_text[0])
    H, W, C = img.shape
    XMin, YMin, XMax, YMax = 0, 0, W - 1, H - 1
    for row in range(H):
        for col in range(W):
            if (img[row, col, :] == bgr).all():
                YMin = row
                break
        else:
            continue
        break
    for row in range(H):
        for col in range(W):
            if (img[H - 1 - row, col, :] == bgr).all():
                YMax = H - 1 - row
                break
        else:
            continue
        break

    for col in range(W):
        for row in range(H):
            if (img[row, col, :] == bgr).all():
                XMin = col
                break
        else:
            continue
        break

    for col in range(W):
        for row in range(H):
            if (img[row, W - 1 - col, :] == bgr).all():
                XMax = W - 1 - col
                break
        else:
            continue
        break
    return XMin, YMin, XMax, YMax


def Find2NumsEdge(img, BackGroundColor):
    H, W, C = img.shape
    res = [[0, 0, 0, 0], [0, 0, 0, 0]]
    for col in range(W):
        if False in (img[:, col, :] == BackGroundColor):
            res[0][0] = col
            break
    for col in range(res[0][0], W):
        if (img[:, col, :] == BackGroundColor).all():
            res[0][2] = col
            break
    for col in range(res[0][2], W):
        if (img[:, col, :] != BackGroundColor).any():
            res[1][0] = col
            break
    for col in range(res[1][0], W):
        if (img[:, col, :] == BackGroundColor).all():
            res[1][2] = col
            break
    for row in range(H):
        if (img[row, res[0][0]: res[0][2], :] != BackGroundColor).any():
            res[0][1] = row
            break
    for row in range(res[0][1], H):
        if (img[row:, res[0][0]: res[0][2], :] == BackGroundColor).all():
            res[0][3] = row
            break
    for row in range(H):
        if (img[row, res[1][0]: res[1][2], :] != BackGroundColor).any():
            res[1][1] = row
            break
    for row in range(res[0][1], H):
        if (img[row:, res[1][0]: res[1][2], :] == BackGroundColor).all():
            res[1][3] = row
            break
    res.append([min(res[0][0], res[1][0]), min(res[0][1], res[1][1]), max(res[0][2], res[1][2]), max(res[0][3], res[1][3])])
    return res


def ReviseInfos(Infos, ReviseMatrix, W, H):
    NewInfos = [Infos[0]]
    for Info in Infos[1:]:
        NewInfo = (np.array(Info[:-1]) + np.array(ReviseMatrix)).tolist()
        if NewInfo[0] < 0:
            NewInfo[0] = 0
        if NewInfo[1] < 0:
            NewInfo[1] = 0
        if NewInfo[2] > W:
            NewInfo[2] = W
        if NewInfo[3] > H:
            NewInfo[3] = H
        NewInfo.append(Info[-1])
        NewInfos.append(NewInfo)
    return NewInfos


def AddNoise(Img, Infos):
    for Info in Infos[1:]:
        StarX = 0.1 * (Info[2] - Info[0]) + Info[0]
        EndX = 0.65 * (Info[2] - Info[0]) + Info[0]
        stary = 8 / 15 * (Info[3] - Info[1]) + Info[1]
        endy = 0.9 * (Info[3] - Info[1]) + Info[1]
        cv2.point


def imread(path):
    return cv2.imdecode(np.fromfile(path, dtype=np.uint8), cv2.IMREAD_UNCHANGED)


def imwrite(path, img):
    cv2.imencode('.jpg', img)[1].tofile(path)  # 保存图片


def bgr2yuv420bin(path_output, img_bgr):
    yuv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2YUV)
    y = yuv[:, :, 0]
    u = yuv[:, :, 1]
    v = yuv[:, :, 2]
    u = cv2.resize(u, (int(u.shape[1] / 2), int(u.shape[0] / 2)))
    v = cv2.resize(v, (int(v.shape[1] / 2), int(v.shape[0] / 2)))
    with open(path_output, 'wb') as f:
        for row in range(y.shape[0]):
            for col in range(y.shape[1]):
                f.write(y[row, col])
        for row in range(u.shape[0]):
            for col in range(u.shape[1]):
                f.write(u[row, col])
                f.write(v[row, col])


def yuv420sp2bgr(path, w, h):
    yuv = np.zeros((h, w, 3), dtype=np.uint8)
    u = np.zeros((int(h / 2), int(w / 2)), dtype=np.uint8)
    v = np.zeros((int(h / 2), int(w / 2)), dtype=np.uint8)
    with open(path, 'rb') as f:
        for row in range(h):
            for col in range(w):
                yuv[row, col, 0] = int(f.read(1).hex(), 16)
        for row in range(int(h / 2)):
            for col in range(int(w / 2)):
                u[row, col] = int(f.read(1).hex(), 16)
                v[row, col] = int(f.read(1).hex(), 16)

    yuv[:, :, 1] = cv2.resize(u, (w, h))
    yuv[:, :, 2] = cv2.resize(v, (w, h))
    bgr = cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR)
    return bgr


def y2gray(path, w, h):
    y = np.zeros((h, w, 1), dtype=np.uint8)
    with open(path, 'rb') as f:
        for row in range(h):
            for col in range(w):
                y[row, col, 0] = int(f.read(1).hex(), 16)
    return y
